/*
 * Javassist, a Java-bytecode translator toolkit.
 * Copyright (C) 1999- Shigeru Chiba. All Rights Reserved.
 *
 * The contents of this file are subject to the Mozilla Public License Version
 * 1.1 (the "License"); you may not use this file except in compliance with
 * the License.  Alternatively, the contents of this file may be used under
 * the terms of the GNU Lesser General Public License Version 2.1 or later,
 * or the Apache License Version 2.0.
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
 * for the specific language governing rights and limitations under the
 * License.
 */

package com.feilong.lib.javassist.bytecode.stackmap;

import java.util.ArrayList;
import java.util.List;

import com.feilong.lib.javassist.ClassPool;
import com.feilong.lib.javassist.NotFoundException;
import com.feilong.lib.javassist.bytecode.BadBytecode;
import com.feilong.lib.javassist.bytecode.ByteArray;
import com.feilong.lib.javassist.bytecode.CodeAttribute;
import com.feilong.lib.javassist.bytecode.ConstPool;
import com.feilong.lib.javassist.bytecode.MethodInfo;
import com.feilong.lib.javassist.bytecode.Opcode;
import com.feilong.lib.javassist.bytecode.StackMap;
import com.feilong.lib.javassist.bytecode.StackMapTable;

/**
 * Stack map maker.
 */
public class MapMaker extends Tracer{

    /**
     * Computes the stack map table of the given method and returns it.
     * It returns null if the given method does not have to have a
     * stack map table or it includes JSR.
     */
    public static StackMapTable make(ClassPool classes,MethodInfo minfo) throws BadBytecode{
        CodeAttribute ca = minfo.getCodeAttribute();
        if (ca == null){
            return null;
        }

        TypedBlock[] blocks;
        try{
            blocks = TypedBlock.makeBlocks(minfo, ca, true);
        }catch (BasicBlock.JsrBytecode e){
            return null;
        }

        if (blocks == null){
            return null;
        }

        MapMaker mm = new MapMaker(classes, minfo, ca);
        try{
            mm.make(blocks, ca.getCode());
        }catch (BadBytecode bb){
            throw new BadBytecode(minfo, bb);
        }

        return mm.toStackMap(blocks);
    }

    /**
     * Computes the stack map table for J2ME.
     * It returns null if the given method does not have to have a
     * stack map table or it includes JSR.
     */
    public static StackMap make2(ClassPool classes,MethodInfo minfo) throws BadBytecode{
        CodeAttribute ca = minfo.getCodeAttribute();
        if (ca == null){
            return null;
        }

        TypedBlock[] blocks;
        try{
            blocks = TypedBlock.makeBlocks(minfo, ca, true);
        }catch (BasicBlock.JsrBytecode e){
            return null;
        }

        if (blocks == null){
            return null;
        }

        MapMaker mm = new MapMaker(classes, minfo, ca);
        try{
            mm.make(blocks, ca.getCode());
        }catch (BadBytecode bb){
            throw new BadBytecode(minfo, bb);
        }
        return mm.toStackMap2(minfo.getConstPool(), blocks);
    }

    public MapMaker(ClassPool classes, MethodInfo minfo, CodeAttribute ca){
        super(classes, minfo.getConstPool(), ca.getMaxStack(), ca.getMaxLocals(), TypedBlock.getRetType(minfo.getDescriptor()));
    }

    protected MapMaker(MapMaker old){
        super(old);
    }

    /**
     * Runs an analyzer (Phase 1 and 2).
     */
    void make(TypedBlock[] blocks,byte[] code) throws BadBytecode{
        make(code, blocks[0]);
        findDeadCatchers(code, blocks);
        try{
            fixTypes(code, blocks);
        }catch (NotFoundException e){
            throw new BadBytecode("failed to resolve types", e);
        }
    }

    // Phase 1

    private void make(byte[] code,TypedBlock tb) throws BadBytecode{
        copyTypeData(tb.stackTop, tb.stackTypes, stackTypes);
        stackTop = tb.stackTop;
        copyTypeData(tb.localsTypes.length, tb.localsTypes, localsTypes);

        traceException(code, tb.toCatch);

        int pos = tb.position;
        int end = pos + tb.length;
        while (pos < end){
            pos += doOpcode(pos, code);
            traceException(code, tb.toCatch);
        }

        if (tb.exit != null){
            for (BasicBlock element : tb.exit){
                TypedBlock e = (TypedBlock) element;
                if (e.alreadySet()){
                    mergeMap(e, true);
                }else{
                    recordStackMap(e);
                    MapMaker maker = new MapMaker(this);
                    maker.make(code, e);
                }
            }
        }
    }

    private void traceException(byte[] code,TypedBlock.Catch handler) throws BadBytecode{
        while (handler != null){
            TypedBlock tb = (TypedBlock) handler.body;
            if (tb.alreadySet()){
                mergeMap(tb, false);
                if (tb.stackTop < 1){
                    throw new BadBytecode("bad catch clause: " + handler.typeIndex);
                }

                tb.stackTypes[0] = merge(toExceptionType(handler.typeIndex), tb.stackTypes[0]);
            }else{
                recordStackMap(tb, handler.typeIndex);
                MapMaker maker = new MapMaker(this);
                maker.make(code, tb);
            }

            handler = handler.next;
        }
    }

    private void mergeMap(TypedBlock dest,boolean mergeStack) throws BadBytecode{
        int n = localsTypes.length;
        for (int i = 0; i < n; i++){
            dest.localsTypes[i] = merge(validateTypeData(localsTypes, n, i), dest.localsTypes[i]);
        }

        if (mergeStack){
            n = stackTop;
            for (int i = 0; i < n; i++){
                dest.stackTypes[i] = merge(stackTypes[i], dest.stackTypes[i]);
            }
        }
    }

    private TypeData merge(TypeData src,TypeData target) throws BadBytecode{
        if (src == target){
            return target;
        }else if (target instanceof TypeData.ClassName || target instanceof TypeData.BasicType){
            return target;
        }else if (target instanceof TypeData.AbsTypeVar){
            ((TypeData.AbsTypeVar) target).merge(src);
            return target;
        }else{
            throw new RuntimeException("fatal: this should never happen");
        }
    }

    private void recordStackMap(TypedBlock target) throws BadBytecode{
        TypeData[] tStackTypes = TypeData.make(stackTypes.length);
        int st = stackTop;
        recordTypeData(st, stackTypes, tStackTypes);
        recordStackMap0(target, st, tStackTypes);
    }

    private void recordStackMap(TypedBlock target,int exceptionType) throws BadBytecode{
        TypeData[] tStackTypes = TypeData.make(stackTypes.length);
        tStackTypes[0] = toExceptionType(exceptionType).join();
        recordStackMap0(target, 1, tStackTypes);
    }

    private TypeData.ClassName toExceptionType(int exceptionType){
        String type;
        if (exceptionType == 0){
            type = "java.lang.Throwable";
        }else{
            type = cpool.getClassInfo(exceptionType);
        }

        return new TypeData.ClassName(type);
    }

    private void recordStackMap0(TypedBlock target,int st,TypeData[] tStackTypes) throws BadBytecode{
        int n = localsTypes.length;
        TypeData[] tLocalsTypes = TypeData.make(n);
        int k = recordTypeData(n, localsTypes, tLocalsTypes);
        target.setStackMap(st, tStackTypes, k, tLocalsTypes);
    }

    protected static int recordTypeData(int n,TypeData[] srcTypes,TypeData[] destTypes){
        int k = -1;
        for (int i = 0; i < n; i++){
            TypeData t = validateTypeData(srcTypes, n, i);
            destTypes[i] = t.join();
            if (t != TOP){
                k = i + 1; // t might be long or double.
            }
        }

        return k + 1;
    }

    protected static void copyTypeData(int n,TypeData[] srcTypes,TypeData[] destTypes){
        for (int i = 0; i < n; i++){
            destTypes[i] = srcTypes[i];
        }
    }

    private static TypeData validateTypeData(TypeData[] data,int length,int index){
        TypeData td = data[index];
        if (td.is2WordType() && index + 1 < length){
            if (data[index + 1] != TOP){
                return TOP;
            }
        }

        return td;
    }

    // Phase 1.5

    /*
     * Javac may generate an exception handler that catches only the exception
     * thrown within the handler itself. It is dead code.
     * See javassist.JvstTest4.testJIRA195().
     */

    private void findDeadCatchers(byte[] code,TypedBlock[] blocks) throws BadBytecode{
        int len = blocks.length;
        for (int i = 0; i < len; i++){
            TypedBlock block = blocks[i];
            if (!block.alreadySet()){
                fixDeadcode(code, block);
                BasicBlock.Catch handler = block.toCatch;
                if (handler != null){
                    TypedBlock tb = (TypedBlock) handler.body;
                    if (!tb.alreadySet()){
                        // tb is a handler that catches only the exceptions
                        // thrown from dead code.
                        recordStackMap(tb, handler.typeIndex);
                        fixDeadcode(code, tb);
                        tb.incoming = 1;
                    }
                }

            }
        }
    }

    private void fixDeadcode(byte[] code,TypedBlock block) throws BadBytecode{
        int pos = block.position;
        int len = block.length - 3;
        if (len < 0){
            // if the dead-code length is shorter than 3 bytes.
            if (len == -1){
                code[pos] = Opcode.NOP;
            }

            code[pos + block.length - 1] = (byte) Opcode.ATHROW;
            block.incoming = 1;
            recordStackMap(block, 0);
            return;
        }

        // if block.incomping > 0, all the incoming edges are from
        // other dead code blocks.  So set block.incoming to 0.
        block.incoming = 0;

        for (int k = 0; k < len; k++){
            code[pos + k] = Opcode.NOP;
        }

        code[pos + len] = (byte) Opcode.GOTO;
        ByteArray.write16bit(-len, code, pos + len + 1);
    }

    // Phase 2

    /*
     * This method first finds strongly connected components (SCCs)
     * in a TypeData graph by Tarjan's algorithm.
     * SCCs are TypeData nodes sharing the same type.
     * Since SCCs are found in the topologically sorted order,
     * their types are also fixed when they are found.
     */
    private void fixTypes(byte[] code,TypedBlock[] blocks) throws NotFoundException,BadBytecode{
        List<TypeData> preOrder = new ArrayList<>();
        int len = blocks.length;
        int index = 0;
        for (int i = 0; i < len; i++){
            TypedBlock block = blocks[i];
            if (block.alreadySet()){ // if block is not dead code
                int n = block.localsTypes.length;
                for (int j = 0; j < n; j++){
                    index = block.localsTypes[j].dfs(preOrder, index, classPool);
                }

                n = block.stackTop;
                for (int j = 0; j < n; j++){
                    index = block.stackTypes[j].dfs(preOrder, index, classPool);
                }
            }
        }
    }

    // Phase 3

    public StackMapTable toStackMap(TypedBlock[] blocks){
        StackMapTable.Writer writer = new StackMapTable.Writer(32);
        int n = blocks.length;
        TypedBlock prev = blocks[0];
        int offsetDelta = prev.length;
        if (prev.incoming > 0){ // the first instruction is a branch target.
            writer.sameFrame(0);
            offsetDelta--;
        }

        for (int i = 1; i < n; i++){
            TypedBlock bb = blocks[i];
            if (isTarget(bb, blocks[i - 1])){
                bb.resetNumLocals();
                int diffL = stackMapDiff(prev.numLocals, prev.localsTypes, bb.numLocals, bb.localsTypes);
                toStackMapBody(writer, bb, diffL, offsetDelta, prev);
                offsetDelta = bb.length - 1;
                prev = bb;
            }else if (bb.incoming == 0){
                // dead code.
                writer.sameFrame(offsetDelta);
                offsetDelta = bb.length - 1;
            }else{
                offsetDelta += bb.length;
            }
        }

        return writer.toStackMapTable(cpool);
    }

    /**
     * Returns true if cur is a branch target.
     */
    private boolean isTarget(TypedBlock cur,TypedBlock prev){
        int in = cur.incoming;
        if (in > 1){
            return true;
        }else if (in < 1){
            return false;
        }

        return prev.stop;
    }

    private void toStackMapBody(StackMapTable.Writer writer,TypedBlock bb,int diffL,int offsetDelta,TypedBlock prev){
        // if diffL is -100, two TypeData arrays do not share
        // any elements.

        int stackTop = bb.stackTop;
        if (stackTop == 0){
            if (diffL == 0){
                writer.sameFrame(offsetDelta);
                return;
            }else if (0 > diffL && diffL >= -3){
                writer.chopFrame(offsetDelta, -diffL);
                return;
            }else if (0 < diffL && diffL <= 3){
                int[] data = new int[diffL];
                int[] tags = fillStackMap(bb.numLocals - prev.numLocals, prev.numLocals, data, bb.localsTypes);
                writer.appendFrame(offsetDelta, tags, data);
                return;
            }
        }else if (stackTop == 1 && diffL == 0){
            TypeData td = bb.stackTypes[0];
            writer.sameLocals(offsetDelta, td.getTypeTag(), td.getTypeData(cpool));
            return;
        }else if (stackTop == 2 && diffL == 0){
            TypeData td = bb.stackTypes[0];
            if (td.is2WordType()){
                // bb.stackTypes[1] must be TOP.
                writer.sameLocals(offsetDelta, td.getTypeTag(), td.getTypeData(cpool));
                return;
            }
        }

        int[] sdata = new int[stackTop];
        int[] stags = fillStackMap(stackTop, 0, sdata, bb.stackTypes);
        int[] ldata = new int[bb.numLocals];
        int[] ltags = fillStackMap(bb.numLocals, 0, ldata, bb.localsTypes);
        writer.fullFrame(offsetDelta, ltags, ldata, stags, sdata);
    }

    private int[] fillStackMap(int num,int offset,int[] data,TypeData[] types){
        int realNum = diffSize(types, offset, offset + num);
        ConstPool cp = cpool;
        int[] tags = new int[realNum];
        int j = 0;
        for (int i = 0; i < num; i++){
            TypeData td = types[offset + i];
            tags[j] = td.getTypeTag();
            data[j] = td.getTypeData(cp);
            if (td.is2WordType()){
                i++;
            }

            j++;
        }

        return tags;
    }

    private static int stackMapDiff(int oldTdLen,TypeData[] oldTd,int newTdLen,TypeData[] newTd){
        int diff = newTdLen - oldTdLen;
        int len;
        if (diff > 0){
            len = oldTdLen;
        }else{
            len = newTdLen;
        }

        if (stackMapEq(oldTd, newTd, len)){
            if (diff > 0){
                return diffSize(newTd, len, newTdLen);
            }else{
                return -diffSize(oldTd, len, oldTdLen);
            }
        }
        return -100;
    }

    private static boolean stackMapEq(TypeData[] oldTd,TypeData[] newTd,int len){
        for (int i = 0; i < len; i++){
            if (!oldTd[i].eq(newTd[i])){
                return false;
            }
        }

        return true;
    }

    private static int diffSize(TypeData[] types,int offset,int len){
        int num = 0;
        while (offset < len){
            TypeData td = types[offset++];
            num++;
            if (td.is2WordType()){
                offset++;
            }
        }

        return num;
    }

    // Phase 3 for J2ME

    public StackMap toStackMap2(ConstPool cp,TypedBlock[] blocks){
        StackMap.Writer writer = new StackMap.Writer();
        int n = blocks.length; // should be > 0
        boolean[] effective = new boolean[n];
        TypedBlock prev = blocks[0];

        // Is the first instruction a branch target?
        effective[0] = prev.incoming > 0;

        int num = effective[0] ? 1 : 0;
        for (int i = 1; i < n; i++){
            TypedBlock bb = blocks[i];
            if (effective[i] = isTarget(bb, blocks[i - 1])){
                bb.resetNumLocals();
                prev = bb;
                num++;
            }
        }

        if (num == 0){
            return null;
        }

        writer.write16bit(num);
        for (int i = 0; i < n; i++){
            if (effective[i]){
                writeStackFrame(writer, cp, blocks[i].position, blocks[i]);
            }
        }

        return writer.toStackMap(cp);
    }

    private void writeStackFrame(StackMap.Writer writer,ConstPool cp,int offset,TypedBlock tb){
        writer.write16bit(offset);
        writeVerifyTypeInfo(writer, cp, tb.localsTypes, tb.numLocals);
        writeVerifyTypeInfo(writer, cp, tb.stackTypes, tb.stackTop);
    }

    private void writeVerifyTypeInfo(StackMap.Writer writer,ConstPool cp,TypeData[] types,int num){
        int numDWord = 0;
        for (int i = 0; i < num; i++){
            TypeData td = types[i];
            if (td != null && td.is2WordType()){
                numDWord++;
                i++;
            }
        }

        writer.write16bit(num - numDWord);
        for (int i = 0; i < num; i++){
            TypeData td = types[i];
            writer.writeVerifyTypeInfo(td.getTypeTag(), td.getTypeData(cp));
            if (td.is2WordType()){
                i++;
            }
        }
    }
}
